{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "55c8f222-47ee-4e07-afd3-079fe41e5505",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.8.0\n",
      "0.9.0\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "import cv2\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision\n",
    "import numpy as np\n",
    "# from utils.general import non_max_suppression,increment_path\n",
    "# from utils.torch_utils import select_device\n",
    "from pathlib import Path\n",
    "device = torch.device(\"cuda\",0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ed566fdc-3a7d-4edb-8c68-bedc9803c082",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/robotnx/.cache/torch/hub/ultralytics_yolov5_master\n",
      "YOLOv5 🚀 2021-10-15 torch 1.8.0 CUDA:0 (Xavier, 7765.4140625MB)\n",
      "\n",
      "Fusing layers... \n",
      "Model Summary: 213 layers, 7225885 parameters, 0 gradients\n",
      "Adding AutoShape... \n"
     ]
    }
   ],
   "source": [
    "# Initialize\n",
    "# if device.type != 'cpu':\n",
    "#         model(torch.zeros(1, 3, 640, 640).to(device).type_as(next(model.parameters())))\n",
    "# half &= device.type != 'cpu'\n",
    "\n",
    "model = torch.hub.load('ultralytics/yolov5', 'yolov5s')\n",
    "video = cv2.VideoCapture(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b4325a95-95ca-4657-b574-aded0382b30d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "box [tensor(224.43750, device='cuda:0'), tensor(407.10938, device='cuda:0'), tensor(342.98438, device='cuda:0'), tensor(865.68750, device='cuda:0')]\n",
      "conf tensor(0.86963, device='cuda:0')\n",
      "person\n",
      "(224, 407) (342, 865)\n",
      "box [tensor(50.09766, device='cuda:0'), tensor(398.03906, device='cuda:0'), tensor(230.34375, device='cuda:0'), tensor(905.34375, device='cuda:0')]\n",
      "conf tensor(0.85352, device='cuda:0')\n",
      "person\n",
      "(50, 398) (230, 905)\n",
      "box [tensor(669.51562, device='cuda:0'), tensor(400.57031, device='cuda:0'), tensor(809.57812, device='cuda:0'), tensor(880.87500, device='cuda:0')]\n",
      "conf tensor(0.83740, device='cuda:0')\n",
      "person\n",
      "(669, 400) (809, 880)\n",
      "box [tensor(4.00781, device='cuda:0'), tensor(207.35156, device='cuda:0'), tensor(810., device='cuda:0'), tensor(799.87500, device='cuda:0')]\n",
      "conf tensor(0.81250, device='cuda:0')\n",
      "bus\n",
      "(4, 207) (810, 799)\n",
      "box [tensor(0., device='cuda:0'), tensor(561.09375, device='cuda:0'), tensor(80.68359, device='cuda:0'), tensor(881.71875, device='cuda:0')]\n",
      "conf tensor(0.30933, device='cuda:0')\n",
      "person\n",
      "(0, 561) (80, 881)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#for i in range(1000):\n",
    "    #flag, img_rd =video.read()\n",
    "import numpy\n",
    "img_rd = 'yolo5/data/images/bus.jpg'\n",
    "\n",
    "pred = model(img_rd)\n",
    "#image = numpy.array(img_rd, dtype = 'uint8')\n",
    "img_rd = cv2.imread(img_rd)\n",
    "for *box, conf, cls in pred.pred[0]:\n",
    "    print('box',box)\n",
    "    print('conf',conf)\n",
    "    print(pred.names[int(cls)] )\n",
    "    leftP =  (int(box[0].item()),int(box[1].item()))\n",
    "    rightP = (int(box[2].item()),int(box[3].item()))\n",
    "\n",
    "    centerx =  int((leftP[0] + rightP[0] )/ 2 - 10)\n",
    "    centery =  int((leftP[1] + rightP[1] )/ 2)\n",
    "    print(leftP,rightP)\n",
    "\n",
    "    cv2.rectangle(img_rd,leftP,rightP,(0, 255, 0), 2)\n",
    "    string =  pred.names[int(cls)]\n",
    "    cv2.putText(img_rd,string,(centerx, centery), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 2)\n",
    "cv2.imwrite(\"3.jpg\",img_rd)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08030b6c-c89d-40b0-8254-5f249584ef43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply NMS\n",
    "pred = non_max_suppression(pred, conf_thres=0.25, classes=2,iou_thres=0.45, max_det=1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "714b6199-3898-4a71-8a56-5add43c5ec94",
   "metadata": {},
   "outputs": [],
   "source": [
    "#webcam0\n",
    "# for i in range(1000):\n",
    "%matplotlib inline\n",
    "flag, img_rd =video.read()\n",
    "cv2.imwrite(\"1.jpg\",img_rd)\n",
    "plt.imshow(img_rd)\n",
    "pred = model(img_rd)\n",
    "pred.print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49e7bfcb-09ae-4496-b45a-c178055c64de",
   "metadata": {},
   "outputs": [],
   "source": [
    "flag, img_rd =video.read()\n",
    "print(flag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0edf51ae-daf9-42db-8fc9-61735a61ab8f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
